"""
离线空气质量预测服务

基于历史数据的长期空气质量预测模块，主要功能包括:
1. 从数据库加载历史空气质量数据
2. 使用预训练的LSTM模型进行多步预测
3. 保存预测结果为JSON文件
4. 生成预测结果可视化图表
5. 可选地将预测结果发送到Kafka

使用说明:
1. 通过命令行参数指定预测天数(--days)和模型路径(--model)
2. 可选启用Kafka发送预测结果(--kafka)
3. 预测结果保存在predictions目录下
4. 可视化图表以PNG格式保存
"""

import os
import sys
import numpy as np
import pandas as pd
import json
import logging
from datetime import datetime, timedelta
import argparse
from tensorflow.keras.models import load_model, Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from sqlalchemy import create_engine, desc, text
import matplotlib.pyplot as plt

from config.db_config import get_db_url

# 配置日志
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger('offline_prediction')

class OfflinePredictionService:
    """离线预测服务"""
    
    def __init__(self, model_path='models/plots/lstm_model.h5', db_path=None, look_back=24):
        """
        初始化离线预测服务
        
        参数:
        model_path: LSTM模型路径
        db_path: 数据库连接URL
        look_back: 时间窗口大小
        """
        self.model_path = model_path
        self.db_path = db_path if db_path else get_db_url()
        self.look_back = look_back
        
        # 加载模型并修改输出层
        if os.path.exists(model_path):
            # 加载原始模型
            original_model = load_model(model_path, custom_objects={'Adam': Adam})
            
            # 验证模型输入形状
            expected_input_shape = (self.look_back, 4)  # 对应4个核心特征
            if original_model.input_shape[1:] != expected_input_shape:
                logger.error(f"模型输入形状不匹配！期望 {expected_input_shape}，实际 {original_model.input_shape[1:]}")
                raise ValueError("模型输入形状与预处理特征数量不匹配")
            
            # 创建新模型，保持相同结构但修改输出层
            self.model = Sequential()
            for layer in original_model.layers[:-1]:  # 复制除输出层外的所有层
                self.model.add(layer)
            
            # 添加新的输出层，输出4个值
            self.model.add(Dense(4))
            
            # 复制权重(除输出层外)
            for i in range(len(original_model.layers)-1):
                self.model.layers[i].set_weights(original_model.layers[i].get_weights())
            
            # 编译模型
            self.model.compile(
                optimizer=Adam(learning_rate=0.001),
                loss='mean_squared_error',
                metrics=['mae']
            )
                
            logger.info(f"加载预测模型: {model_path}")
        else:
            logger.error(f"模型文件不存在: {model_path}")
            raise FileNotFoundError(f"模型文件不存在: {model_path}")
        
        # 创建数据库连接
        self.engine = create_engine(self.db_path)
        

    def enable_kafka(self):
        """
        启用Kafka发送预测结果
        """
        try:
            from utils.kafka_client import KafkaProducer
            from config.kafka_config import KAFKA_CONFIG
            
            self.kafka_producer = KafkaProducer(
                bootstrap_servers=KAFKA_CONFIG.get('bootstrap_servers', 'localhost:9092'),
                topic=KAFKA_CONFIG.get('prediction_topic', 'air_quality_predictions')
            )
            logger.info(f"已启用Kafka发送预测结果")
        except Exception as e:
            logger.error(f"启用Kafka失败: {str(e)}")
    
    def process_data(self, df):
        """
        处理数据，包括缺失值处理、异常值处理等
        
        参数:
            df (pandas.DataFrame): 原始数据
            
        返回:
            pandas.DataFrame: 处理后的数据
        """
        # 转换时间戳
        df['timestamp'] = pd.to_datetime(df['date'].astype(str) + ' ' + df['hour'].astype(str).str.zfill(2), format='%Y%m%d %H')
        
        # 选择所有污染物指标
        # 使用模型预期的5个核心特征（排除NO2）
        core_features = ['AQI', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3']
        available_features = [col for col in core_features if col in df.columns]
        
        if not available_features:
            logger.error("未找到任何污染物数据列")
            return None
            
        logger.info(f"处理以下污染物数据: {available_features}")
        
        # 初始化结果DataFrame
        result_df = df[['timestamp'] + available_features].copy()
        
        # 按时间排序
        result_df = result_df.sort_values('timestamp')
        
        # 对所有特征进行缺失值处理
        for col in available_features:
            result_df[col] = result_df[col].fillna(method='ffill').fillna(method='bfill')
            
            # 检查是否还有缺失值
            if result_df[col].isna().any():
                logger.warning(f"{col} 列仍存在缺失值，使用0填充")
                result_df[col] = result_df[col].fillna(0)
        
        # 数据标准化/归一化应在预测函数中进行
        # 这样可以确保与训练数据使用相同的缩放参数
        
        return result_df
    
    def load_historical_data(self, days=30):
        """
        从数据库或CSV文件加载历史数据
        
        参数:
            days (int): 要加载的历史数据天数，默认30天
            
        返回:
            DataFrame: 处理后的空气质量数据，包含时间戳和污染物数值
            None: 如果加载失败
            
        注意:
            1. 优先从数据库加载数据
            2. 如果数据库加载失败，尝试从CSV文件加载
            3. 自动填充缺失值和异常值
        """
        try:
            # 计算开始日期
            end_date = datetime.now()
            start_date = end_date - timedelta(days=days)
            
            try:
                # 尝试从数据库加载数据
                query = text("""
                    SELECT date, hour, type, site_id, value 
                    FROM air_quality 
                    WHERE CONCAT(date, LPAD(hour,2,'0')) >= DATE_FORMAT(:start_date, '%Y%m%d%H')
        AND CONCAT(date, LPAD(hour,2,'0')) <= DATE_FORMAT(:end_date, '%Y%m%d%H')
                    ORDER BY date, hour
                """)
                
                df = pd.read_sql_query(query, self.engine, params={'start_date': start_date.date(), 'end_date': end_date.date()})
                
                if not df.empty:
                    logger.info(f"从数据库加载了{len(df)}条历史数据")
                    processed_df = self.process_data(df)
                    return processed_df
            except Exception as e:
                logger.warning(f"从数据库加载数据失败，尝试从CSV文件加载: {str(e)}")
            
            # 如果数据库加载失败，尝试从CSV文件加载
            csv_dir = os.path.join('data', 'result')
            if not os.path.exists(csv_dir):
                logger.error(f"CSV目录不存在: {csv_dir}")
                return None
                
            # 获取所有CSV文件并按日期排序
            csv_files = [f for f in os.listdir(csv_dir) if f.endswith('.csv') and f.startswith('china_sites_')]
            csv_files.sort()
            
            # 读取并合并CSV文件
            dfs = []
            for csv_file in csv_files[-days:]:
                try:
                    file_path = os.path.join(csv_dir, csv_file)
                    df = pd.read_csv(file_path)
                    dfs.append(df)
                except Exception as e:
                    logger.warning(f"读取CSV文件失败 {csv_file}: {str(e)}")
                    continue
            
            if dfs:
                combined_df = pd.concat(dfs, ignore_index=True)
                logger.info(f"从CSV文件加载了{len(combined_df)}条历史数据")
                processed_df = self.process_data(combined_df)
                return processed_df
            else:
                logger.warning(f"未找到{days}天内的历史数据")
                return None
                
        except Exception as e:
            logger.error(f"加载历史数据失败: {str(e)}")
            return None
    
    def prepare_prediction_data(self, df):
        """
        准备预测数据
        
        参数:
        df: 处理后的DataFrame
        
        返回:
        X: 模型输入数据
        scaler: 标准化器
        """
        from sklearn.preprocessing import MinMaxScaler
        
        try:
            # 选择所有污染物指标
            features = ['AQI', 'PM2.5', 'PM10', 'SO2', 'NO2', 'O3']
            available_features = [col for col in features if col in df.columns]
            
            # 严格使用核心特征(与模型训练保持一致)
            required_features = ['PM2.5', 'PM10', 'SO2', 'O3']  # 使用4个特征
            missing_features = [f for f in required_features if f not in available_features]
            if missing_features:
                logger.error(f"缺失必要特征: {missing_features}")
                return None, None
                
            # 确保使用与模型训练相同的特征顺序
            available_features = [f for f in required_features if f in available_features]
            
            if not available_features:
                logger.error("未找到任何污染物数据列")
                return None, None
            
            # 确保数据按时间排序
            df = df.sort_values('timestamp')
            
            # 检查数据完整性
            for col in available_features:
                if df[col].isnull().any():
                    logger.error(f"{col} 列存在缺失值")
                    return None, None
                    
            if len(df) < self.look_back:
                logger.error(f"数据量不足，需要至少{self.look_back}条记录")
                return None, None
            
            # 打印可用特征和数量用于调试
            logger.info(f"可用特征: {available_features}, 数量: {len(available_features)}")
            
            # 获取最后look_back条记录
            recent_data = df[available_features].values[-self.look_back:]
            
            # 确保数据形状为(1, look_back, 4)
            if recent_data.shape[1] != 4:
                logger.error(f"数据特征数量不匹配，需要4个特征，实际有{recent_data.shape[1]}个")
                logger.info(f"当前特征: {available_features}")
                # 填充缺失特征为0
                if recent_data.shape[1] < 4:
                    padding = np.zeros((recent_data.shape[0], 4 - recent_data.shape[1]))
                    recent_data = np.hstack((recent_data, padding))
                    logger.warning(f"填充缺失特征为0，调整后特征数量: {recent_data.shape[1]}")
                    logger.info(f"填充前的形状: {df[available_features].values[-self.look_back:].shape}")
                    logger.info(f"填充后的形状: {recent_data.shape}")
                else:
                    return None, None
            
            # 验证最终输入形状
            logger.info(f"最终模型输入形状: {recent_data.reshape(1, self.look_back, len(available_features)).shape}")
            
            # 检查数据范围
            if np.any(np.isinf(recent_data)) or np.any(np.isnan(recent_data)):
                logger.error("数据包含无效值(inf或nan)")
                return None, None
            
            # 检查数据类型
            if not np.issubdtype(recent_data.dtype, np.number):
                logger.error(f"数据类型错误，期望数值类型，实际类型: {recent_data.dtype}")
                return None, None
            
            # 标准化数据
            scaler = MinMaxScaler(feature_range=(0, 1))
            # 使用所有数据来拟合scaler以获得更好的缩放范围
            scaler.fit(df[available_features].values)
            
            # 将最后look_back条记录标准化
            normalized_data = scaler.transform(recent_data)
            
            # 检查标准化后的数据
            if np.any(np.isinf(normalized_data)) or np.any(np.isnan(normalized_data)):
                logger.error("标准化后的数据包含无效值")
                return None, None
            
            # 调整为模型输入形状 [1, look_back, num_features]
            X = normalized_data.reshape(1, self.look_back, len(available_features))
            
            logger.info(f"成功准备预测数据，输入形状: {X.shape}")
            return X, scaler
            
        except Exception as e:
            logger.error(f"准备预测数据失败: {str(e)}")
            return None, None
    
    def predict_future(self, days=7):
        """
        预测未来空气质量
        
        参数:
            days (int): 预测天数，默认7天
            
        返回:
            list: 预测结果列表，每个元素包含:
                - timestamp: 预测时间点
                - 各污染物预测值(PM2_5, PM10等)
                - prediction_type: 预测类型(offline)
                - hours_ahead: 预测提前小时数
            None: 如果预测失败
            
        算法说明:
            1. 使用滑动窗口方法进行多步预测
            2. 每次预测一个时间点，然后更新输入窗口
            3. 使用MinMaxScaler进行数据标准化/反标准化
            4. 自动生成预测时间戳序列
        """
        try:
            # 加载历史数据
            df = self.load_historical_data(days=max(30, days+3))  # 确保有足够的历史数据
            if df is None:
                logger.error("无法加载历史数据，预测失败")
                return None
            
            # 准备预测数据
            X, scaler = self.prepare_prediction_data(df)
            if X is None or scaler is None:
                logger.error("准备预测数据失败")
                return None
            
            # 预测未来几天
            hours_to_predict = days * 24
            
            # 使用最近数据作为起点
            current_window = X[0]
            predictions = []
            
            for i in range(hours_to_predict):
                # 预测下一个时间点
                next_pred = self.model.predict(current_window[np.newaxis, :, :], verbose=0)
                predictions.append(next_pred[0])
                
                # 更新窗口（滑动窗口）
                current_window = np.vstack([current_window[1:], next_pred])
            
            # 反标准化预测结果
            denormalized_predictions = scaler.inverse_transform(np.array(predictions))
            
            # 确保所有污染物值为非负数
            denormalized_predictions = np.maximum(denormalized_predictions, 0)
            
            # 生成时间戳
            last_timestamp = datetime.now().replace(minute=0, second=0, microsecond=0) + timedelta(hours=1)
            timestamps = [last_timestamp + timedelta(hours=i) for i in range(hours_to_predict)]
            
            # 构建结果
            results = []
            for i, (ts, pred) in enumerate(zip(timestamps, denormalized_predictions)):
                result = {
                    'timestamp': ts.isoformat(),
                    'prediction_type': 'offline',
                    'hours_ahead': i + 1
                }
                
                # 添加所有污染物预测值(使用与模型训练相同的4个特征)
                features = ['PM2.5', 'PM10', 'SO2', 'O3']
                for j, col in enumerate(features):
                    if j < len(pred):  # 确保不越界
                        result[col] = max(0, float(pred[j]))  # 双重确保非负
                
                results.append(result)
                

            
            logger.info(f"完成了{days}天的离线预测，共{len(results)}个小时")
            
            # 保存预测结果
            self._save_predictions(results, days)
            
            return results
            
        except Exception as e:
            logger.error(f"预测失败: {str(e)}")
            return None
    
    def _save_predictions(self, predictions, days):
        """
        保存预测结果
        
        参数:
        predictions: 预测结果列表
        days: 预测天数
        """
        try:
            # 确保目录存在
            os.makedirs('predictions', exist_ok=True)
            
            # 保存为JSON
            timestamp = datetime.now().strftime('%Y%m%d%H%M')
            filename = f"predictions/offline_{days}days_{timestamp}.json"
            
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(predictions, f, indent=4, ensure_ascii=False)
            
            logger.info(f"预测结果已保存至: {filename}")
            
            # 创建可视化
            self._visualize_predictions(predictions, days)
            
        except Exception as e:
            logger.error(f"保存预测结果失败: {str(e)}")
    
    def _visualize_predictions(self, predictions, days):
        """
        可视化预测结果
        
        参数:
        predictions: 预测结果列表
        days: 预测天数
        """
        try:
            # 转换为DataFrame
            df = pd.DataFrame(predictions)
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            df.set_index('timestamp', inplace=True)
            
            # 设置图表大小
            plt.figure(figsize=(15, 10))
            
            # 绘制预测的污染物(只绘制实际预测的4个特征)
            pollutants = ['PM2.5', 'PM10', 'SO2', 'O3']
            colors = ['blue', 'green', 'red', 'purple']
            
            for pollutant, color in zip(pollutants, colors):
                if pollutant in df.columns:
                    plt.plot(df.index, df[pollutant], label=pollutant, color=color)
            
            # 添加图例和标题
            plt.legend()
            plt.title(f'未来{days}天空气质量预测（离线模式）')
            plt.xlabel('时间')
            plt.ylabel('污染物浓度')
            plt.grid(True)
            
            # 如果预测天数大于7天，使用周为单位标记x轴
            if days > 7:
                from matplotlib.dates import WeekdayLocator, DateFormatter
                plt.gca().xaxis.set_major_locator(WeekdayLocator())
                plt.gca().xaxis.set_major_formatter(DateFormatter('%m-%d'))
            
            # 保存图表
            timestamp = datetime.now().strftime('%Y%m%d%H%M')
            filename = f"predictions/offline_{days}days_{timestamp}.png"
            plt.savefig(filename)
            plt.close()
            
            logger.info(f"预测结果可视化已保存至: {filename}")
            
        except Exception as e:
            logger.error(f"可视化预测结果失败: {str(e)}")
    
    def close(self):
        """关闭资源"""
        pass

def main():
    parser = argparse.ArgumentParser(description='离线空气质量预测服务')
    parser.add_argument('--days', type=int, default=7, help='预测天数（默认7天）')
    parser.add_argument('--model', type=str, default='models/plots/lstm_model.h5', help='模型路径')
    parser.add_argument('--kafka', action='store_true', help='启用Kafka发送预测结果')
    
    args = parser.parse_args()
    
    try:
        # 检查模型是否存在
        if not os.path.exists(args.model):
            logger.error(f"模型文件不存在: {args.model}")
            return
        
        # 创建离线预测服务
        service = OfflinePredictionService(model_path=args.model)
        
        # 如果启用Kafka，初始化生产者
        if args.kafka:
            service.enable_kafka()
        
        # 执行预测
        predictions = service.predict_future(days=args.days)
        
        if predictions:
            logger.info(f"离线预测完成，共生成{len(predictions)}个预测数据点")
        else:
            logger.error("离线预测失败")
        
    except KeyboardInterrupt:
        logger.info("收到终止信号")
    except Exception as e:
        logger.error(f"服务运行时发生错误: {str(e)}")
    finally:
        if 'service' in locals():
            service.close()

if __name__ == "__main__":
    main()
